#!/usr/bin/env bash

# STRICT MODE
set -euo pipefail
#set -vx

# COLORS

BLACK='\e[0;30m'
DARK_GRAY='\e[1;30m'
RED='\e[0;31m'
LIGHT_RED='\e[1;31m'
GREEN='\e[0;32m'
LIGHT_GREEN='\e[1;32m'
BROWN_ORANGE='\e[0;33m'
YELLOW='\e[1;33m'
BLUE='\e[0;34m'
LIGHT_BLUE='\e[1;34m'
PURPLE='\e[0;35m'
LIGHT_PURPLE='\e[1;35m'
CYAN='\e[0;36m'
LIGHT_CYAN='\e[1;36m'
LIGHT_GRAY='\e[0;37m'
WHITE='\e[1;37m'

NC='\033[0m'

# Обрабатываем опции командной строки
OUTPUT_TO_FILE=false
CATEGORIZE=false
CATEGORIES_ONLY=false
WORK_TYPE=not_set
SINGLE_TASK_MODE=false
SINGLE_ZNO_MODE=false
VERBOSE=':'
START_DATE=''
END_DATE=''
OUTPUT='stdout'

# Если запущено без параметров вообще, то выводим сразу помощь
[ $# -eq 0 ] && { cat help.resource; exit 0;}

while [ "$#" -gt 0 ]; do
  case "$1" in
    --zno ) WORK_TYPE=zno; shift 1;;
    --zpi ) WORK_TYPE=zpi; shift 1;;
    -c | --categorize      ) CATEGORIZE=true; shift 1;;
    -C | --categories-only ) CATEGORIZE=true; CATEGORIES_ONLY=true; WORK_TYPE=zpi; shift 1;;
    -e | --end-date        ) END_DATE=$2; shift 2;;
    -F | --to-file         ) OUTPUT_TO_FILE=true; shift 1;;
    -h | --help            ) cat help.resource ; exit 0;;
    -I | --incident        ) SINGLE_TASK_MODE=true ; INCIDENT_FROM_CMD_LINE=$2; WORK_TYPE=zpi; shift 2;;
    -O | --output          ) OUTPUT=$2; shift 2;;
    -s | --start-date      ) START_DATE=$2; shift 2;;
    -V | --verbose         ) VERBOSE=echo; shift 1;;
    -Z                     ) SINGLE_ZNO_MODE=true; ZNO_FROM_CMD_LINE=$2; WORK_TYPE=zno; shift 2;;
     * ) echo Что-то пошло не так при парсинге аргументов. Не продолжаем. ; exit 1;;
  esac
done

function get_dates {
  # Нужно обязательно проверить, что переменная не пустая, иначе date вернёт текущее время
#  $VERBOSE Начальная дата: `date -d "$START_DATE" +%c`
  [[ $START_DATE != '' ]] && START_DATE=`date -d "$START_DATE" +%s` || return 3
#  $VERBOSE Начальная дата: `date -d "$START_DATE" +%c`

# TODO Подумать над отключением проверок, по факту работают только при отключении строгого режима
  [ $? -ne 0 ] && { echo Неверный формат начальной даты; exit 2 ; }

#  $VERBOSE Конечная дата: `date -d "$END_DATE" +%c`
  END_DATE=`date -d "$END_DATE" +%s`
  # TODO Подумать над отключением проверки
  [ $? -ne 0 ] && { echo Неверный формат конечной даты; exit 2 ; }

  return 0 # Без этого не работает в строгом режиме
}

# Проверки опций
[[ $WORK_TYPE == 'not_set' ]] && { echo Необходимо обязательно указать опции либо --zpi, либо --zno. Не продолжаем.; exit 1; }
[[ $CATEGORIZE == 'true' ]] && [[ ! $WORK_TYPE == 'zpi' ]] && { echo Опция -c может быть использована только совместно с --zpi. Не продолжаем.; exit 1; }
[[ $START_DATE != '' && $END_DATE == '' ]] && echo Не указано конечное время. Будет использоваться текущее время в качестве конечного
[[ $START_DATE == '' && $END_DATE != '' ]] && { echo Указано конечное время без начального. Не продолжаем.; exit 1;}
# TODO проверка аргумента в -I, первые два символа чтобы были IM или первые четыре IMTS

# Глобальные переменные
read THESAURUS <<< `cat thesaurus | tr '\n' ' '`
TMP_COUNTER=0
DIFF_COUNTER=0

function tag_is_correct \
{
  for TAG_FROM_LIST in ${TAGS_ARRAY[*]} ; do [[ $TAG_FROM_LIST == $TAG ]] && { return 0; break ;} ; done
  return 1
}

function process_zno \
{
  (( CURRENT_ITERATION > ZNO_TRIES )) && { echo Дошли до 4 итерации на ЗНО $ID, тэг $TAG >> errlog ; TAG=NO_TAG ; return 0; }
  (( CURRENT_ITERATION = CURRENT_ITERATION + 1 ))

  # Забираем заголовок и тело ЗНО из базы
  REQUEST_SUBJECT=`psql -d sm_helper -tAc "SELECT translate(subject, '\\\''\"', ' ') FROM zno_uploaded WHERE number = '$ID'"    | tr '\n' ' '`
  REQUEST_DESCRIPTION=`psql -d sm_helper -tAc "SELECT translate(description, '\\\''\"', ' ') FROM zno_uploaded WHERE number = '$ID'" | tr '\n' ' '`

  $VERBOSE REQUEST_SUBJECT $REQUEST_SUBJECT
  $VERBOSE REQUEST_DESCRIPTION $REQUEST_DESCRIPTION
  $VERBOSE Обрабатываем ЗНО $ID

  # Основное действие, запрос к нейросети
  TAG=`gs "Напиши тэг. Твоя задача определить правильный тэг для запроса пользователя. В ответе должен быть только тэг, никаких других пояснений. Вот список возможных тэгов: $REQUEST_TAGS. Вот словарь сокращений: $THESAURUS Вот запрос пользователя: $REQUEST_SUBJECT. $REQUEST_DESCRIPTION"`
#  TAG_2=`gs "Проанализируй текст, и выбери для него тэг. В ответе должен быть только тэг, никаких других пояснений. В тексте могут использоваться следующие со

  $VERBOSE -----------------------
  $VERBOSE $TAG
  #$VERBOSE $TAG_2
  $VERBOSE -----------------------

  tag_is_correct || process_zno

  return 0
}

function zno_output \
{
  # Закидываем тэг в базу
  [[ $OUTPUT == 'database' ]] && psql --quiet -d sm_helper -c "UPDATE zno_uploaded SET tag = '$TAG' WHERE number = '$ID'"
#  [[ $OUTPUT == 'stdout' ]] && echo -e Тема ЗНО: $REQUEST_SUBJECT \\n Описание ЗНО: $REQUEST_DESCRIPTION \\n Тэг номер один: $TAG \\n Тэг номер два: $TAG_2 \$
  [[ $OUTPUT == 'stdout' ]] && { echo Тема ЗНО: $REQUEST_SUBJECT ; echo Описание ЗНО: $REQUEST_DESCRIPTION ; echo Тэг номер один: $TAG ; }
#  [[ ! $TAG == $TAG_2 ]] && { echo -e $ID\\n $REQUEST_DESCRIPTION\\n $TAG\\n $TAG_2\\n >> diff_tags_2 ; let "DIFF_COUNTER = DIFF_COUNTER + 1"; }
  # echo -e К ЗНО \\n $REQUEST_SUBJECT $REQUEST_DESCRIPTION \\nГигачат не придумал нормального тэга. \\nВот такой тэг: $TAG

  return 0
}

function do_znos {
  read REQUEST_TAGS <<< `cat /home/postgres/shevchenko/gigachat/sm_helper/tag_list | tr '\n' ' '`
#  REQUEST_TAGS=`cat $TAGFILE`
  ZNO_TMP_FILE=/home/postgres/shevchenko/gigachat/sm_helper/zno.tmp
  rm $ZNO_TMP_FILE
  TAGS_ARRAY=(PG_EXTENSION PG_REPLICATION PG_DUMP PG_ROLE PG_HBA PG_RLM PG_REPAIR PG_ANALYZE PG_BOUNCER PG_SSL PG_CONFIG PG_CONSULTATION PG_OTHER)
  ZNO_TRIES=3

  # Переключение между режимом одного ЗНО и обработкой всех ЗНО
  if [[ $SINGLE_ZNO_MODE == 'true' ]] ; then
    echo $ZNO_FROM_CMD_LINE > $ZNO_TMP_FILE
    TOTAL_ZNO_COUNT=1
  else
    psql -tAc "SELECT number FROM zno_uploaded WHERE type = 'ЗНО'" -d sm_helper -o $ZNO_TMP_FILE
    TOTAL_ZNO_COUNT=`psql -tAc "SELECT count(number) FROM zno_uploaded WHERE type = 'ЗНО'" -d sm_helper`
  fi

  for ID in `cat $ZNO_TMP_FILE` ; do
    CURRENT_ITERATION=0

    process_zno

    # Если гигачат не придумал нормальный тэг, то ставим NO_TAG
    # tag_is_correct || TAG=NO_TAG

    $VERBOSE после проверок $TAG

    zno_output

    let "TMP_COUNTER = TMP_COUNTER + 1"
    echo -e ${BLUE}Обработано ЗНО:${NC} $TMP_COUNTER ${BLUE}/ $TOTAL_ZNO_COUNT${NC}
  done
}

function do_incidents {
    CATEGORY_FILE=/home/postgres/shevchenko/gigachat/sm_helper/category_list
    REQUEST_TAGS=`cat $CATEGORY_FILE`

    # TODO Вынести это в отдельную функцию
    # Забираем спиcок инцидентов из базы
    if [[ $SINGLE_TASK_MODE == 'false' ]] ; then
      if get_dates ; then
        declare -a INCIDENTS=`psql -tAc "SELECT incident_number FROM incidents WHERE extract( epoch from start_time ) >= '$START_DATE' AND extract( epoch from end_time ) <= '$END_DATE'" -d sm_helper | tr '\n' ' '`
      else
        declare -a INCIDENTS=`psql -tAc "select incident_number from incidents" -d sm_helper | tr '\n' ' '`
      fi
    else
        declare -a INCIDENTS=($INCIDENT_FROM_CMD_LINE)
    fi

    $VERBOSE Список инцидентов, по которым работаем: ${INCIDENTS[*]}

    # Обрабатываем каждый инцидент
    for CURRENT_INCIDENT in ${INCIDENTS[*]} ; do
        declare -a ZPIS
        # Забираем номера всех ЗПИ по текущему инциденту
        ZPIS=`psql -tAc "SELECT zpi_number FROM zpis WHERE parent = '$CURRENT_INCIDENT'" -d sm_helper | tr '\n' ' '`

        # Собираем текст из всех ЗПИ, и соединяем в одну строку
        ZPI_TEXT=''
        for CURRENT_ZPI in ${ZPIS[*]} ; do
            ZPI_TEXT="$ZPI_TEXT`psql -tAc "SELECT description FROM zpis WHERE zpi_number = '$CURRENT_ZPI'" -d sm_helper | tr '\n' ' '` "
        done

        # Забираем текст самого инцидента
        INCIDENT_TEXT=`psql -tAc "SELECT description FROM incidents WHERE incident_number = '$CURRENT_INCIDENT'" -d sm_helper | tr '\n' ' '`

        # Если включен флаг -c, то категоризируем по тэгам
        [[ $CATEGORIZE == 'true' ]] && TAG_CATEGORY=`gs "Напиши тэг. Твоя задача определить правильный тэг для запроса пользователя. В ответе должен быть только тэг, никаких других пояснений. Вот список возможных тэгов: $REQUEST_TAGS. Вот словарь сокращений: $THESAURUS Вот запрос пользователя: $INCIDENT_TEXT. $ZPIS"`

        # Получаем ответ от нейросети
        TAG_DESCRIPTIVE=`gs "Напиши одной фразой короткое описание к инциденту пользователя. Не указывай IP-адресов и имён хостов. Не придумывай ничего. Если данных слишком мало, просто напиши not_enough_data. Вот словарь используемых сокращений: $THESAURUS Вот инцидент пользователя: $INCIDENT_TEXT. $ZPIS"`
        TAG_ANALYTICAL=`gs "Проанализируй проблему с точки зрения базы данных, СУБД, ответь в терминах специалистов баз данных. Дай краткую выжимку, одной фразой, без имён хостов и дополнительной информации. Не придумывай ничего. $ADDITIONAL_PROMPT Если данных слишком мало, просто напиши not_enough_data. Вот словарь используемых сокращений: $THESAURUS Вот инцидент пользователя: $INCIDENT_TEXT. Вот запросы по данному инциденту: $ZPIS"`

        # Выводим инфу
        if [[ $OUTPUT_TO_FILE == 'false' ]] ; then
            if [[ $CATEGORIES_ONLY == 'false' ]] ; then
                echo -e Текст инцидента $CURRENT_INCIDENT:\\n$INCIDENT_TEXT\\n\\n
                echo -e Текст всех ЗПИ вместе:\\n$ZPI_TEXT\\n\\n
                echo -e ${BLUE}Описательный ответ Гигачата:${NC}\\n$TAG_DESCRIPTIVE
                echo -e ${LIGHT_BLUE}Аналитичный ответ Гигачата:${NC}\\n$TAG_ANALYTICAL
                [[ $CATEGORIZE == 'true' ]] && echo -e ${RED}Категория инцидента:${NC} $TAG_CATEGORY
            else
                echo -e ${LIGHT_BLUE}Номер инцидента: $CURRENT_INCIDENT${NC}\\n${BLUE}Категория: $TAG_CATEGORY${NC}
            fi
        else
            echo -e Описательный ответ Гигачата:\\n$TAG_DESCRIPTIVE >> answer_0${TMP_COUNTER}_desc
            echo -e Аналитичный ответ Гигачата:\\n$TAG_ANALYTICAL  >> answer_0${TMP_COUNTER}_anal
        fi
    done
}


function main {
  [[ $WORK_TYPE == 'zno' ]] && do_znos
  [[ $WORK_TYPE == 'zpi' ]] && do_incidents
}

main
